import cv2
import numpy as np


# 将灰度值量化为 2 位（4 级）
def quantize_to_2bit(gray_frame):
    quantized_frame = np.zeros_like(gray_frame, dtype=np.uint8)
    quantized_frame[gray_frame < 64] = 0  # 0-63 -> 00 (二进制)
    quantized_frame[(gray_frame >= 64) & (gray_frame < 128)] = 1  # 64-127 -> 01 (二进制)
    quantized_frame[(gray_frame >= 128) & (gray_frame < 192)] = 2  # 128-191 -> 10 (二进制)
    quantized_frame[gray_frame >= 192] = 3  # 192-255 -> 11 (二进制)
    return quantized_frame


# 将量化后的 2 位灰度值压缩存储，每个字节存储 4 个像素
def compress_2bit_frame(quantized_frame):
    height, width = quantized_frame.shape
    compressed_size = (height * width) // 4  # 每 4 个像素压缩成 1 字节
    compressed_frame = np.zeros(compressed_size, dtype=np.uint8)

    idx = 0
    for i in range(0, height * width, 4):
        compressed_byte = (
                (quantized_frame.flat[i] << 6) |
                (quantized_frame.flat[i + 1] << 4) |
                (quantized_frame.flat[i + 2] << 2) |
                (quantized_frame.flat[i + 3])
        )
        compressed_frame[idx] = compressed_byte
        idx += 1

    return compressed_frame


# 解压 2 位灰度值，将压缩后的数据还原为 2 位灰度图
def decompress_2bit_frame(compressed_frame, height, width):
    decompressed_frame = np.zeros((height, width), dtype=np.uint8)

    idx = 0
    for i in range(0, height * width, 4):
        compressed_byte = compressed_frame[idx]
        decompressed_frame.flat[i] = (compressed_byte >> 6) & 0b11
        decompressed_frame.flat[i + 1] = (compressed_byte >> 4) & 0b11
        decompressed_frame.flat[i + 2] = (compressed_byte >> 2) & 0b11
        decompressed_frame.flat[i + 3] = compressed_byte & 0b11
        idx += 1

    return decompressed_frame


# 显示解压后的 2 位灰度图，扩展到 8 位灰度图以便于显示
def display_2bit_frame(decompressed_frame):
    expanded_frame = decompressed_frame * 85  # 将 2 位灰度值扩展到 8 位
    return expanded_frame


# 处理视频，将其转换为16x16的2位灰度视频，并输出压缩后的内容
def process_video(video_path, output_file, output_video_path):
    cap = cv2.VideoCapture(video_path)

    # 设置输出视频分辨率为 16x16
    output_width, output_height = 16, 16
    fps = cap.get(cv2.CAP_PROP_FPS)  # 获取原视频的帧率

    # 打开视频写入器，保存解压后的 16x16 灰度视频
    fourcc = cv2.VideoWriter_fourcc(*'XVID')
    out = cv2.VideoWriter(output_video_path, fourcc, fps, (output_width, output_height), isColor=False)

    frame_number = 0
    with open(output_file, 'w') as f:
        f.write("// Compressed 16x16 frames stored in 2-bit grayscale\n\n")

        while cap.isOpened():
            ret, frame = cap.read()
            if not ret:
                break

            # 将帧转换为灰度图
            gray_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

            # 将灰度图缩放为 16x16
            small_gray_frame = cv2.resize(gray_frame, (output_width, output_height), interpolation=cv2.INTER_AREA)

            # 将灰度图量化为 2 位
            quantized_frame = quantize_to_2bit(small_gray_frame)

            # 压缩 2 位灰度图
            compressed_frame = compress_2bit_frame(quantized_frame)

            # 输出压缩后的内容到文件
            f.write(f"// Frame {frame_number}\n")
            f.write(f"const uint8_t Frame{frame_number}[] PROGMEM = {{\n")
            for i in range(0, len(compressed_frame), 16):
                line = ",".join(f"0x{byte:02X}" for byte in compressed_frame[i:i + 16])
                f.write(f"  {line},\n")
            f.write("};\n\n")

            # 解压 2 位灰度图
            decompressed_frame = decompress_2bit_frame(compressed_frame, output_height, output_width)

            # 展示解压后的图像
            expanded_frame = display_2bit_frame(decompressed_frame)
            expanded_frame_expanded=cv2.resize(expanded_frame, (160, 160))
            cv2.imshow('16x16 2-bit Grayscale Video Preview', expanded_frame_expanded)

            # 保存解压后的灰度视频
            out.write(expanded_frame)

            frame_number += 1

            # 按下 'q' 键退出
            if cv2.waitKey(int(1000 / fps)) & 0xFF == ord('q'):
                break

    cap.release()
    out.release()
    cv2.destroyAllWindows()


# 主程序
video_path = '2.mp4'  # 替换为你的视频文件路径
output_file = 'compressed_frames_output.txt'  # 保存压缩内容的文件
output_video_path = 'output_16x16_2bit_video.avi'  # 保存解压后的视频文件

# 处理视频，将其转换为 16x16 的 2 位灰度视频，并输出压缩后的内容
process_video(video_path, output_file, output_video_path)

print(f"视频处理完毕，压缩后的内容已保存到 {output_file}，转换后的 16x16 2 位灰度视频已保存为 {output_video_path}")